"""Script to run analyses for the project."""
import argparse
import os
import re
import sys

import matplotlib.pyplot as plt
import pandas as pd

from cat_rfs.code.estimate_betas import estimate_betas
from cat_rfs.code.fig2 import print_fig2  # Fig 1 is printed with print_table3
from cat_rfs.code.fig3 import print_fig3
from cat_rfs.code.fig4 import print_fig4
from cat_rfs.code.fig5 import print_fig5
from cat_rfs.code.table1 import print_table1
from cat_rfs.code.table2 import print_table2
from cat_rfs.code.table3 import print_table3
from cat_rfs.code.table4 import print_table4
from cat_rfs.code.table5 import print_table5
from cat_rfs.code.table6 import print_table6


def main(output='console', reconstruct_betas=False):
    """Main entry point for the analysis."""
    assert output in ('console', 'paper')
    assert reconstruct_betas in (True, False)

    plt.ion()  # Turn on interactive mode for figures

    PERILS = ['Asia Earthquake', 'Asia Cyclone',
              'Europe Earthquake', 'Europe Cyclone',
              'North_america Earthquake', 'North_america Cyclone',
              'Middle_america Earthquake', 'Middle_america Cyclone',
              'South_america Earthquake']

    # %% Load data, apply filters, and estimate betas (if reconstruct_betas)
    # or load pre-estimated betas (if not reconstruct_betas)

    # Load data
    df = pd.read_csv('cat_rfs/data/primary.csv', parse_dates=['maturity', 'pricing_date'])
    df = df.drop(columns=['size', 'spread', 'attachment_prob', 'expected_loss', 'exhaustion_prob',
                          'attachment_prob_wsst', 'expected_loss_wsst', 'exhaustion_prob_wsst'])

    df2 = pd.read_csv('cat_rfs/data/secondary.csv', parse_dates=['date', 'current_maturity'])
    df = df.merge(df2, on=['CUSIP9'], validate='1:m')

    # Filters
    # 1. Sample start
    df = df[df['date'] >= '6/30/2003']

    # 2. Drop distressed bonds
    df = df[df['distress_fl'] == 0]  # "Distress" flag in secondary market pricing sheet of Lane
    df = df[df['rating_change'] == 0]  # No credit rating agency has changed bond rating since issuance
    df = df[df['creditwatch'] == 0]  # Bond is not under credit watch by any agency
    df = df[df['triggered'] == 0]  # Bond has not triggered

    # 3. Bond is not maturing within a year
    df = df.loc[((df['current_maturity'] - df['date'])) > pd.Timedelta('273d')]

    # 4. No missing observations for key variables
    df = df.dropna(subset=['attachment_prob', 'expected_loss', 'exhaustion_prob'])

    cols = [f"el_{c.replace(' ', '_').lower()}" for c in PERILS]
    df = df.dropna(subset=cols, how='all')
    df[cols] = df[cols].fillna(0)  # Value is empty if bond has no exposure to particular peril. Fill with zeros.

    cols = [f"el_{c.replace(' ', '_').lower()}_wsst" for c in PERILS]
    df = df.dropna(subset=cols, how='all')
    df[cols] = df[cols].fillna(0)  # Value is empty if bond has no exposure to particular peril. Fill with zeros.

    # 5. Bond is exposed to at least one peril for which exposure data exists
    df = df[df['perils'].str.contains('|'.join(map(re.escape, PERILS)), case=False, na=False)].copy()

    # Get betas
    if reconstruct_betas:
        estimate_betas(df)

    df = df.merge(pd.read_csv(f"cat_rfs/data/betas{'_reconstructed' if reconstruct_betas else ''}.csv",
                              parse_dates=['date']), on=['CUSIP9', 'date'], how='inner', validate='1:1')

    # %% Print summary stats
    print_table1(output=output, reconstruct_betas=reconstruct_betas)

    # %% Run analyses
    print_table2(df, output=output)
    print_table3(df, output=output)
    print_table4(df, output=output)
    print_table5(df, output=output)
    print_table6(df, output=output)

    print_fig2(output=output)
    print_fig3(df, output=output)
    print_fig4(output=output)
    print_fig5(df, output=output, reconstruct_variances=reconstruct_betas)

    pass


if __name__ == '__main__':
    # Change working directory to the parent folder of the script
    script_dir = os.path.dirname(os.path.abspath(__file__))
    parent_dir = os.path.dirname(script_dir)
    os.chdir(parent_dir)

    # Parse optional arguments
    parser = argparse.ArgumentParser(description='Run main script with optional arguments.')
    parser.add_argument('--output', type=str, default='console', help='Output mode (default: console, alternative: paper)')
    parser.add_argument('--reconstruct_betas', action='store_true', help='Enable beta reconstruction')
    args = parser.parse_args()

    # Run main
    sys.exit(main(output=args.output, reconstruct_betas=args.reconstruct_betas))
